package com.mrgood.config;

import org.springframework.stereotype.Component;
import org.springframework.web.servlet.handler.HandlerInterceptorAdapter;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.*;

@Component
public class TracingInterceptor extends HandlerInterceptorAdapter {
    public final static ThreadLocal<Map<String, Collection<String>>> headerThreadLocal = new ThreadLocal<>();

    private static final List<String> tracingHeaderKeys = Arrays.asList("x-request-id", "x-b3-traceid", "x-b3-spanid", "x-b3-parentspanid",
            "x-b3-sampled", "x-b3-flags", "x-ot-span-context","User-Agent","cookie");

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) {
        Map<String, Collection<String>> tracingHeaders = new HashMap<>();
        for (String key : tracingHeaderKeys) {
            String val = request.getHeader(key);
            if (null != val) {
                tracingHeaders.put(key, Arrays.asList(val));
            }
        }
        if (tracingHeaders.size() > 0) {
            headerThreadLocal.set(tracingHeaders);
        }
        response.setHeader("traceid", request.getHeader("x-b3-traceid"));
        return true;
    }

}
